"""
MongoDB初始化配置模块
负责数据库初始化、用户认证设置和索引创建
"""
import logging
from pymongo import MongoClient, ASCENDING
from pymongo.errors import ServerSelectionTimeoutError, OperationFailure
from typing import Optional

class MongoDBInitializer:
    def __init__(self, host: str = 'localhost', port: int = 27017):
        self.host = host
        self.port = port
        self.logger = logging.getLogger(__name__)
        self.client = None

    def initialize_database(self, 
                          db_name: str, 
                          username: str, 
                          password: str,
                          auth_source: str = 'admin') -> bool:
        """
        初始化MongoDB数据库
        
        Args:
            db_name: 数据库名称
            username: 管理员用户名
            password: 管理员密码
            auth_source: 认证数据库
            
        Returns:
            bool: 初始化是否成功
        """
        try:
            # 1. 连接MongoDB
            self.client = MongoClient(self.host, self.port, serverSelectionTimeoutMS=5000)
            
            # 2. 检查连接
            self.client.admin.command('ping')
            
            # 3. 创建管理员用户（如果不存在）
            self._create_admin_user(username, password)
            
            # 4. 使用认证重新连接
            self.client = MongoClient(
                host=self.host,
                port=self.port,
                username=username,
                password=password,
                authSource=auth_source
            )
            
            # 5. 创建数据库和集合
            db = self.client[db_name]
            self._create_collections(db)
            
            # 6. 创建索引
            self._create_indexes(db)
            
            self.logger.info(f"MongoDB初始化成功: {db_name}")
            return True
            
        except ServerSelectionTimeoutError as e:
            self.logger.error(f"MongoDB连接失败: {str(e)}")
            return False
        except Exception as e:
            self.logger.error(f"MongoDB初始化失败: {str(e)}")
            return False
        
    def _create_admin_user(self, username: str, password: str) -> None:
        """创建管理员用户"""
        try:
            self.client.admin.command(
                "createUser",
                username,
                pwd=password,
                roles=[
                    {"role": "userAdminAnyDatabase", "db": "admin"},
                    {"role": "readWriteAnyDatabase", "db": "admin"}
                ]
            )
            self.logger.info(f"创建管理员用户成功: {username}")
        except OperationFailure as e:
            if "already exists" in str(e):
                self.logger.info(f"管理员用户已存在: {username}")
            else:
                raise

    def _create_collections(self, db) -> None:
        """创建必要的集合"""
        # 提交历史集合
        if "commits" not in db.list_collection_names():
            db.create_collection("commits")
            
        # 文件历史集合
        if "files" not in db.list_collection_names():
            db.create_collection("files")
            
        # 合并请求集合
        if "merge_requests" not in db.list_collection_names():
            db.create_collection("merge_requests")

    def _create_indexes(self, db) -> None:
        """创建索引"""
        # commits集合索引
        db.commits.create_index([("commit_id", ASCENDING)], unique=True)
        db.commits.create_index([("date", ASCENDING)])
        db.commits.create_index([("author", ASCENDING)])
        
        # files集合索引
        db.files.create_index([("file_path", ASCENDING)])
        db.files.create_index([("language", ASCENDING)])
        
        # merge_requests集合索引
        db.merge_requests.create_index([("id", ASCENDING)], unique=True)
        db.merge_requests.create_index([("created_at", ASCENDING)])

    def check_connection(self, username: str, password: str, db_name: str) -> bool:
        """测试数据库连接"""
        try:
            client = MongoClient(
                host=self.host,
                port=self.port,
                username=username,
                password=password,
                authSource='admin'
            )
            client[db_name].command('ping')
            return True
        except Exception as e:
            self.logger.error(f"连接测试失败: {str(e)}")
            return False 